"""
SEG2格式写入模块
实现自定义的SEG2格式写入功能，因为ObsPy不支持写入SEG2格式

SEG2格式规范参考:
- SEG2是一种简单的地震数据格式
- 包含文件头和道头信息
- 数据通常为16位或32位整数
"""

import os
import struct
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional, Union
from obspy import Stream, Trace


class SEG2Writer:
    """SEG2格式写入器"""
    
    def __init__(self, stream: Stream):
        self.stream = stream
        
    def write(self, filename: Union[str, Path]) -> None:
        """
        将ObsPy Stream写入SEG2格式文件
        
        参数:
            filename: 输出文件名
        """
        filename = Path(filename)
        
        with open(filename, 'wb') as f:
            # 写入文件头
            self._write_file_header(f)
            
            # 写入道数据
            for i, trace in enumerate(self.stream):
                self._write_trace_header(f, trace, i)
                self._write_trace_data(f, trace)
    
    def _write_file_header(self, f) -> None:
        """写入SEG2文件头"""
        # SEG2文件头格式
        # 文件标识符 (4字节)
        f.write(b'SEG2')
        
        # 文件头大小 (2字节)
        file_header_size = 32  # 基本文件头大小
        f.write(struct.pack('<H', file_header_size))
        
        # 道数 (2字节)
        num_traces = len(self.stream)
        f.write(struct.pack('<H', num_traces))
        
        # 采样率 (4字节，浮点数)
        if len(self.stream) > 0:
            sampling_rate = 1.0 / self.stream[0].stats.delta
        else:
            sampling_rate = 100.0  # 默认采样率
        f.write(struct.pack('<f', sampling_rate))
        
        # 数据格式 (2字节)
        # 1 = 16位整数, 2 = 32位整数, 3 = 32位浮点数
        data_format = 1  # 默认使用16位整数
        f.write(struct.pack('<H', data_format))
        
        # 保留字段 (18字节)
        f.write(b'\x00' * 18)
    
    def _write_trace_header(self, f, trace: Trace, trace_number: int) -> None:
        """写入道头"""
        # 道标识符 (4字节)
        f.write(b'TRC2')
        
        # 道头大小 (2字节)
        trace_header_size = 64  # 基本道头大小
        f.write(struct.pack('<H', trace_header_size))
        
        # 道序号 (2字节)
        f.write(struct.pack('<H', trace_number + 1))
        
        # 采样点数 (4字节)
        npts = trace.stats.npts
        f.write(struct.pack('<I', npts))
        
        # 采样间隔 (4字节，微秒)
        delta_microseconds = int(trace.stats.delta * 1e6)
        f.write(struct.pack('<I', delta_microseconds))
        
        # 开始时间 (8字节，Unix时间戳，微秒)
        start_time = int(trace.stats.starttime.timestamp * 1e6)
        f.write(struct.pack('<Q', start_time))
        
        # 台站代码 (8字节)
        station = trace.stats.get('station', 'STA')
        station_bytes = station.ljust(8, '\x00')[:8].encode('ascii')
        f.write(station_bytes)
        
        # 分量代码 (4字节)
        channel = trace.stats.get('channel', 'Z')
        channel_bytes = channel.ljust(4, '\x00')[:4].encode('ascii')
        f.write(channel_bytes)
        
        # 保留字段 (28字节)
        f.write(b'\x00' * 28)
    
    def _write_trace_data(self, f, trace: Trace) -> None:
        """写入道数据"""
        # 将数据转换为16位整数
        data = trace.data
        
        # 归一化数据到16位范围
        if data.dtype == np.float32 or data.dtype == np.float64:
            # 浮点数数据，归一化到16位范围
            max_val = np.max(np.abs(data))
            if max_val > 0:
                scale_factor = 32767.0 / max_val
                data = (data * scale_factor).astype(np.int16)
            else:
                data = np.zeros_like(data, dtype=np.int16)
        else:
            # 整数数据，直接转换
            data = data.astype(np.int16)
        
        # 写入数据
        f.write(data.tobytes())


def write_seg2(stream: Stream, filename: Union[str, Path]) -> None:
    """
    将ObsPy Stream写入SEG2格式文件
    
    参数:
        stream: ObsPy Stream对象
        filename: 输出文件名
    """
    writer = SEG2Writer(stream)
    writer.write(filename)


def convert_segy_to_seg2_custom(segy_file: Union[str, Path], 
                               output_file: Optional[Union[str, Path]] = None) -> Path:
    """
    使用自定义方法将SEGY文件转换为SEG2格式
    
    参数:
        segy_file: SEGY文件路径
        output_file: 输出文件路径（可选）
        
    返回:
        生成的SEG2文件路径
    """
    segy_path = Path(segy_file)
    
    if not segy_path.exists():
        raise FileNotFoundError(f"SEGY文件不存在: {segy_path}")
    
    # 确定输出文件路径
    if output_file is not None:
        seg2_path = Path(output_file)
    else:
        seg2_path = segy_path.parent / f"{segy_path.stem}.seg2"
    
    # 读取SEGY文件
    try:
        st = read(str(segy_path), format="SEGY")
        
        if len(st) == 0:
            raise ValueError(f"SEGY文件不包含任何道: {segy_path}")
        
        # 使用自定义方法写入SEG2
        write_seg2(st, seg2_path)
        
        print(f"已成功将SEGY文件转换为SEG2格式: {seg2_path}")
        return seg2_path
        
    except Exception as e:
        raise RuntimeError(f"转换SEGY到SEG2失败: {e}")


def convert_stream_to_seg2(stream: Stream, 
                          output_dir: Union[str, Path],
                          filename_prefix: str = "data") -> Path:
    """
    将ObsPy Stream转换为SEG2格式
    
    参数:
        stream: ObsPy Stream对象
        output_dir: 输出目录
        filename_prefix: 文件名前缀
        
    返回:
        生成的SEG2文件路径
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 生成文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{filename_prefix}_{timestamp}.seg2"
    seg2_path = output_dir / filename
    
    # 写入SEG2文件
    write_seg2(stream, seg2_path)
    
    return seg2_path


if __name__ == "__main__":
    # 测试代码
    import sys
    from obspy import read
    
    if len(sys.argv) < 2:
        print("用法: python seg2_writer.py <SEGY文件> [输出文件]")
        sys.exit(1)
    
    segy_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None
    
    try:
        seg2_file = convert_segy_to_seg2_custom(segy_file, output_file)
        print(f"转换成功: {seg2_file}")
    except Exception as e:
        print(f"转换失败: {e}")
        sys.exit(1)